
import sys
import os
import io
import argparse
import json
import fitz
from openpyxl import Workbook
from openpyxl.styles import Border, Side
# 解决print 中文乱码问题
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')

# 获取命令行参数
parser = argparse.ArgumentParser()
# 文件数组
parser.add_argument('-filePath', type=str, default='[]', help='argument 1')
# 目录配置信息
parser.add_argument('-outputManage', type=str, default='{}', help='argument 2')

args = parser.parse_args()


def main_2(path, outputManage):
    # 获取路径配置信息、文件列表
    equalSource = json.loads(outputManage)['equalSource']
    outputPath = json.loads(outputManage)['outputPath']
    filePaths = json.loads(path)
    
    # 遍历文件路径数组
    for file in filePaths:
        # 获取文件名称
        file_name = os.path.splitext(file['fileName'])[0]
        # pdf 源文件目录
        pdf_parent_path = os.path.dirname(file['filePath'])
        docx_name = ''
        # 当equalSource为false时，转换的doc文件统一放在配置的目录中
        if equalSource == False:
            # docx文件名称，equalSource为false时
            docx_name = outputPath + '\\' + file_name
        else:
            # docx文件保存在源文件目录下
            docx_name = pdf_parent_path + '\\' + file_name
            
        # 先创建对应的pdf存放目录
        # folder_name = "new_folder"
        if not os.path.exists(docx_name):
            os.mkdir(docx_name)

        # 读取pdf文件
        doc = fitz.open(file['filePath'])
        # text：create a text output
        out = open(docx_name + '\\' + "text.txt", "wb")
        
        # excel： 创建工作簿对象
        workbook = Workbook()
        # excel 选取默认的活动表格（第一个）
        sheet = workbook.active
        # excel 定义边框样式
        border_style = Border(left=Side(style='thin'), right=Side(style='thin'), top=Side(style='thin'), bottom=Side(style='thin'))
        
        # 循环读取每页pdf
        for page_index in range(len(doc)):  # iterate over pdf pages
            page = doc[page_index]  # get the page
            # 读取图片
            image_list = page.get_images()
            # 读取text
            text = page.get_text().encode("utf8")  # get plain text (is in UTF-8)
            # 读取 table
            tabs = page.find_tables()  # locate and extract any tables on page
            # 提取excel print(tabs)
            if len(tabs.tables):
                print(f"{len(tabs.tables)} found on {page}")
                
                # print(tabs.tables)
                for table_index in range(len(tabs.tables)):
                    table_one = tabs[table_index].extract()

                    # print(f"{len(table_one)} 长度。")
                    # 遍历数组并写入单元格
                    for row in table_one:
                        sheet.append(row)
                        
                # 遍历每个单元格并应用边框样式
                for row in sheet.iter_rows():
                    for cell in row:
                        cell.border = border_style
                    # sheet[row].row_dimensions[row].height = 20
                # 设置行高为20
                for row in range(1, sheet.max_row + 1):
                    sheet.row_dimensions[row].height = 20

                # 设置列宽为15
                for column in range(1, sheet.max_column + 1):
                    sheet.column_dimensions[chr(64+column)].width = 15
                # 保存工作簿到指定路径
                workbook.save(docx_name + '\\' + "excel.xlsx")

            # print the number of images found on the page
            # if image_list:
            #     print(f"Found {len(image_list)} images on page {page_index}")
            # else:
            #     print("No images found on page", page_index)

            # 提取纯文本
            out.write(text)  # write text of page
            out.write(bytes((12,)))  # write page delimiter (form feed 0x0C) 页面分隔符
            
            # 提取图片 enumerate the image list
            for image_index, img in enumerate(image_list, start=1):
                xref = img[0]  # get the XREF of the image
                pix = fitz.Pixmap(doc, xref)  # create a Pixmap

                if pix.n - pix.alpha > 3:  # CMYK: convert to RGB first
                    pix = fitz.Pixmap(fitz.csRGB, pix)

                # save the image as png
                pix.save(docx_name + '\\' + ("page_%s-image_%s.png" % (page_index, image_index)))
                pix = None
         
        # 关闭text提取        
        out.close()
        print(file['fileId'])

if __name__ == '__main__':
    main_2(args.filePath, args.outputManage)
